import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
import numpy as np

import warnings

warnings.filterwarnings('ignore')

MIN_GROWTH_RATE = 0.05
MIN_VAL = 1  #
medal_counts_df = pd.read_csv('D:\Desktop\MCM_Python\成品部分_代码_lunw/初步数据/summerOly_medal_counts_cleaned.csv',encoding='latin1')
athletes_df = pd.read_csv('D:\Desktop\MCM_Python\成品部分_代码_lunw/初步数据/summerOly_athletes.csv', encoding='latin1')

medal_counts_df['Year'] = medal_counts_df['Year'].astype(int)
athletes_df['Year'] = athletes_df['Year'].astype(int)
athletes_count = athletes_df.groupby(['NOC', 'Year']).size().reset_index(name='Athletes')
merged_df = pd.merge(medal_counts_df, athletes_count, on=['NOC', 'Year'], how='left')
events_count = athletes_df.groupby(['NOC', 'Year'])['Sport'].nunique().reset_index(name='Events')
merged_df = pd.merge(merged_df, events_count, on=['NOC', 'Year'], how='left')

# 2024年数据
athletes_2024 = merged_df[merged_df['Year'] == 2024].groupby('NOC')['Athletes'].mean()
events_2024 = merged_df[merged_df['Year'] == 2024].groupby('NOC')['Events'].mean()

athletes_2024 = athletes_2024[athletes_2024.index != 'Mixed team']
events_2024 = events_2024[events_2024.index != 'Mixed team']

athletes_2024.fillna(MIN_VAL, inplace=True)
athletes_2024[athletes_2024 < MIN_VAL] = MIN_VAL
events_2024.fillna(MIN_VAL,inplace=True)
events_2024[events_2024 < MIN_VAL] = MIN_VAL


def calculate_growth_rate(df, column):
    # 分组后，如果只有一条记录或首条=0，返回MIN_GROWTH_RATE
    def _growth(group):
        if len(group) > 1 and group[column].iloc[0] != 0:
            ratio = (group[column].iloc[-1] / group[column].iloc[0]) ** (1 / (len(group) - 1)) - 1
            return ratio if ratio > 0 else MIN_GROWTH_RATE
        else:
            return MIN_GROWTH_RATE

    result = df.groupby('NOC').apply(_growth).reset_index(name=f'{column}_growth_rate')
    return result


# 计算增长率
athletes_growth = calculate_growth_rate(merged_df, 'Athletes')
events_growth = calculate_growth_rate(merged_df, 'Events')

athletes_2024_growth = athletes_growth.set_index('NOC')['Athletes_growth_rate']
events_2024_growth = events_growth.set_index('NOC')['Events_growth_rate']

# 通过计算年增长率来预测2028年的值
athletes_2028 = athletes_2024 * (1 + athletes_2024_growth)
events_2028 = events_2024 * (1 + events_2024_growth)

predictions_gold = {}
predictions_total = {}
predictions_interval_gold = {}
predictions_interval_total = {}

for country in merged_df['NOC'].unique():
    country_data = merged_df[merged_df['NOC'] == country]
    if len(country_data) < 2:
        print(f"Skipping {country}, insufficient data.")
        continue

    X = country_data[['Year', 'Athletes', 'Events']]
    y_gold = country_data['Gold']
    y_total = country_data['Total']
    X_train, X_test, y_train_gold, y_test_gold = train_test_split(X, y_gold, test_size=0.2, random_state=42)
    X_train_total, X_test_total, y_train_total, y_test_total = train_test_split(X, y_total, test_size=0.2,
                                                                                random_state=42)

    model_gold = RandomForestRegressor(n_estimators=300, random_state=42)
    model_total = RandomForestRegressor(n_estimators=300, random_state=42)
    model_gold.fit(X_train, y_train_gold)
    model_total.fit(X_train_total, y_train_total)
    mse_gold = mean_squared_error(y_test_gold, model_gold.predict(X_test))
    mse_total = mean_squared_error(y_test_total, model_total.predict(X_test_total))

    print(f"{country}: Gold MSE={mse_gold}, Total MSE={mse_total}")

    if country in athletes_2028.index and country in events_2028.index:
        X_2028 = pd.DataFrame({
            'Year': [2028],
            'Athletes': [athletes_2028[country]],
            'Events': [events_2028[country]]
        })
        X_2028 = X_2028.replace([np.inf, -np.inf], 0).fillna(0)

        pred_gold_2028 = model_gold.predict(X_2028)
        pred_total_2028 = model_total.predict(X_2028)

        # 不确定性区间
        interval_g = np.std([model_gold.predict(X_2028) for _ in range(100)])
        interval_t = np.std([model_total.predict(X_2028) for _ in range(100)])

        predictions_gold[country] = pred_gold_2028[0]
        predictions_total[country] = pred_total_2028[0]
        predictions_interval_gold[country] = interval_g
        predictions_interval_total[country] = interval_t

# 2024年未得牌的国家
no_medals_2024 = merged_df[
    (merged_df['Year'] == 2024)
    & (merged_df['Gold'] == 0)
    & (merged_df['Silver'] == 0)
    & (merged_df['Bronze'] == 0)
    ]



predictions_first_medal = {}
for country in no_medals_2024['NOC'].unique():
    # 如果这个国家没有计算出增长率,或者不在index中,跳过
    if country not in athletes_2028.index or country not in events_2028.index:
        predictions_first_medal[country] = 0
        continue

    # 构建2028
    X_2028 = pd.DataFrame({
        'Year': [2028],
        'Athletes': [athletes_2028[country]],
        'Events': [events_2028[country]]
    })
    X_2028 = X_2028.replace([np.inf, -np.inf], 0).fillna(0)
    pred_g = model_gold.predict(X_2028)[0]
    predictions_first_medal[country] = 1 if pred_g > 0 else 0

count_first_medal = sum(predictions_first_medal.values())
total_first_medal = len(predictions_first_medal)
probability = count_first_medal / total_first_medal if total_first_medal > 0 else 0

print(f"{count_first_medal} countries might get first medal in 2028.")
print(f"Probability= {probability * 100:.2f}%")

first_medal_df = pd.DataFrame({
    'NOC': list(predictions_first_medal.keys()),
    'Predicted First Medal in 2028': list(predictions_first_medal.values())
})
first_medal_df.to_csv('first_medal_predictions.csv', index=False)

predictions_df = pd.DataFrame({
    'NOC': list(predictions_gold.keys()),
    'Predicted Gold Medals for 2028': list(predictions_gold.values()),
    'Predicted Total Medals for 2028': list(predictions_total.values()),
    'Gold Prediction Interval': list(predictions_interval_gold.values()),
    'Total Prediction Interval': list(predictions_interval_total.values())
})
predictions_df.to_csv('2028_predictions_by_country.csv', index=False)